import argparse
from random import shuffle
from utils import *

try:
    from FPMC_numba import FPMC
except ImportError:
    from FPMC import FPMC
# from FPMC import FPMC

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input_dir', help='The directory of input', type=str)
    parser.add_argument('-e', '--n_epoch', help='# of epoch', type=int, default=10)
    parser.add_argument('--n_neg', help='# of neg samples', type=int, default=20)
    parser.add_argument('-n', '--n_factor', help='dimension of factorization', type=int, default=32)
    parser.add_argument('-l', '--learn_rate', help='learning rate', type=float, default=0.01)
    parser.add_argument('-r', '--regular', help='regularization', type=float, default=0.01)
    args = parser.parse_args()

    f_dir = args.input_dir

    data_list, user_set, ques_set = load_data_from_dir(f_dir)
    shuffle(data_list)

    train_ratio = 0.8
    split_idx = int(len(data_list) * train_ratio)
    tr_data = data_list[:split_idx]
    te_data = data_list[split_idx:]

    fpmc = FPMC(n_user=max(user_set) + 1, n_item=max(ques_set) + 1,
                n_factor=args.n_factor, learn_rate=args.learn_rate, regular=args.regular)
    fpmc.user_set = user_set
    fpmc.item_set = ques_set
    fpmc.init_model()

    precision_in, precision_out, recall_in, recall_out, HR_in, HR_out, mrr_in, mrr_out = \
        (fpmc.learnSBPR_FPMC(tr_data, te_data=te_data, n_epoch=args.n_epoch, neg_batch_size=args.n_neg,
                             eval_per_epoch=True))
    print('Precision (In/Out): %.4f/%.4f, Recall (In/Out): %.4f/%.4f, HR (In/Out): %.4f/%.4f, MRR (In/Out): %.4f/%.4f'
          % (precision_in, precision_out, recall_in, recall_out, HR_in, HR_out, mrr_in, mrr_out))
    # precision_out, recall_out, auc_out, mrr_out = \
    #     (fpmc.learnSBPR_FPMC(tr_data, te_data=te_data, n_epoch=args.n_epoch, neg_batch_size=args.n_neg,
    #                          eval_per_epoch=False))
    # print('Precision: %.4f, Recall: %.4f, AUC: %.4f, MRR: %.4f' % (precision_out, recall_out, auc_out, mrr_out))
